from __future__ import division
from math import sqrt as sqrt
from itertools import product as product
import torch


class PriorBox(object):
    """Compute priorbox coordinates in center-offset form for each source
    feature map.
    """
    def __init__(self, cfg):
        super(PriorBox, self).__init__()
        self.image_size = cfg['min_dim']
        # number of priors for feature map location (either 4 or 6)
        self.num_priors = len(cfg['aspect_ratios'])
        self.variance = cfg['variance'] or [0.1]
        self.feature_maps = cfg['feature_maps'] # feature_map：此处代表了每个特征图的大小
        self.min_sizes = cfg['min_sizes']   # min_size：此参数为每个抽取的特征图上的一个点对应原图上的框的最小值
        self.max_sizes = cfg['max_sizes']   # max_size：此参数实际上为min_size[k+1]
        self.steps = cfg['steps']   # steps：此处代表了此特征图相对原图的下采样倍数
        self.aspect_ratios = cfg['aspect_ratios']
        self.clip = cfg['clip'] # clip：控制字，控制是否要将坐标压缩在[0,1]以内，以防越界
        self.version = cfg['name']
        for v in self.variance:
            if v <= 0:
                raise ValueError('Variances must be greater than 0')
    
    def forward(self):
        mean = []   # 用于存放预选框的列表
        for k, f in enumerate(self.feature_maps):
            for i, j in product(range(f), repeat=2):    # 对特征图上每一个点(i,j)进行密集采样
                f_k = self.image_size / self.steps[k]    # f_k计算的是对应特征图的尺寸（如果能整除的话数值上应该等于参数feature_map的大小）
                # 以下两行与源代码不同，用于测试源代码是否有错误
                cx = (i + 0.5) / f_k    # 计算特征图每个方格中心点横坐标对于原图归一化后的坐标
                cy = (j + 0.5) / f_k    # 计算特征图每个方格中心点纵坐标对于原图归一化后的坐标

                # 计算长宽比为1的小预选框尺寸，边长为参数min_size
                s_k = self.min_sizes[k] / self.image_size
                mean += [cx, cy, s_k, s_k]

                # 计算长宽比为1的大预选框尺寸，归一化后边长为参数sqrt(s_k * s_(k+1))
                if self.max_sizes:
                    s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
                    mean += [cx, cy, s_k_prime, s_k_prime]

                # 计算其他自定义长宽比的预选框
                for ar in self.aspect_ratios[k]:
                    mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                    mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]

        # 将结果转换为tensor，按n行4列的方式组织，每一行就是一个预选框对于归一化的（横坐标，纵坐标，宽，高）
        output = torch.Tensor(mean).view(-1, 4)
        if self.clip:
            output.clamp_(min=0, max=1)
        return output

                 


# if __name__ == "__main__":
#     p_b = PriorBox(prior_box_config)
#     output = p_b.forward()
#     print(output.size())
